/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.serialization.protobuf;

import com.google.protobuf.BoolValue;
import com.google.protobuf.BytesValue;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.DoubleValue;
import com.google.protobuf.Empty;
import com.google.protobuf.ExtensionRegistryLite;
import com.google.protobuf.FloatValue;
import com.google.protobuf.GeneratedMessageV3.Builder;
import com.google.protobuf.Int32Value;
import com.google.protobuf.Int64Value;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.MessageLite;
import com.google.protobuf.MessageOrBuilder;
import com.google.protobuf.Parser;
import com.google.protobuf.StringValue;
import com.google.protobuf.util.JsonFormat;
import com.google.protobuf.util.JsonFormat.Printer;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class ProtobufUtils {

  static boolean isSupported(Class<?> clazz) {
    if (clazz == null) {
      return false;
    }

    if (MessageLite.class.isAssignableFrom(clazz)) {
      return true;
    }
    return false;
  }

  /* Protobuf json */

  static <T> T deserializeJson(String json, Class<T> requestClass)
      throws InvalidProtocolBufferException {
    Builder builder;
    try {
      builder = getMessageBuilder(requestClass);
    } catch (Exception e) {
      throw new IllegalArgumentException(
          "Get google protobuf message builder from " + requestClass.getName() + "failed", e);
    }
    JsonFormat.parser().merge(json, builder);
    return (T) builder.build();
  }

  static String serializeJson(Object value) throws InvalidProtocolBufferException {
    Printer printer = JsonFormat.printer().omittingInsignificantWhitespace();
    return printer.print((MessageOrBuilder) value);
  }

  private static Builder getMessageBuilder(Class<?> requestType) throws Exception {
    Method method = requestType.getMethod("newBuilder");
    return (Builder) method.invoke(null, null);
  }

  /* Protobuf */
  private static ConcurrentMap<Class<? extends MessageLite>, MessageMarshaller> marshallers =
      new ConcurrentHashMap<>();

  private static volatile ExtensionRegistryLite globalRegistry =
      ExtensionRegistryLite.getEmptyRegistry();

  static {
    // Built-in types need to be registered in advance
    marshaller(Empty.getDefaultInstance());
    marshaller(BoolValue.getDefaultInstance());
    marshaller(Int32Value.getDefaultInstance());
    marshaller(Int64Value.getDefaultInstance());
    marshaller(FloatValue.getDefaultInstance());
    marshaller(DoubleValue.getDefaultInstance());
    marshaller(BytesValue.getDefaultInstance());
    marshaller(StringValue.getDefaultInstance());
  }

  public static <T extends MessageLite> void marshaller(T defaultInstance) {
    marshallers.put(defaultInstance.getClass(), new MessageMarshaller<>(defaultInstance));
  }

  static void serialize(Object value, OutputStream os) throws IOException {
    MessageLite messageLite = (MessageLite) value;
    messageLite.writeDelimitedTo(os);
  }

  @SuppressWarnings("unchecked")
  static <T> T deserialize(InputStream is, Class<T> requestClass)
      throws InvalidProtocolBufferException {
    MessageMarshaller<?> marshaller = marshallers.get(requestClass);
    if (marshaller == null) {
      throw new IllegalStateException(
          String.format(
              "Protobuf classes should be registered in advance before "
                  + "do serialization, class name: %s",
              requestClass.getName()));
    }
    return (T) marshaller.parse(is);
  }

  public static Empty convertNullToEmpty() {
    return Empty.newBuilder().build();
  }

  public static Object convertEmptyToNull(Empty empty) {
    return null;
  }

  private static final class MessageMarshaller<T extends MessageLite> {
    private final Parser<T> parser;
    private final T defaultInstance;

    @SuppressWarnings("unchecked")
    MessageMarshaller(T defaultInstance) {
      this.defaultInstance = defaultInstance;
      parser = (Parser<T>) defaultInstance.getParserForType();
    }

    @SuppressWarnings("unchecked")
    public Class<T> getMessageClass() {
      return (Class<T>) defaultInstance.getClass();
    }

    public T getMessagePrototype() {
      return defaultInstance;
    }

    public T parse(InputStream stream) throws InvalidProtocolBufferException {
      return parser.parseDelimitedFrom(stream, globalRegistry);
    }

    private T parseFrom(CodedInputStream stream) throws InvalidProtocolBufferException {
      T message = parser.parseFrom(stream, globalRegistry);
      try {
        stream.checkLastTagWas(0);
        return message;
      } catch (InvalidProtocolBufferException e) {
        e.setUnfinishedMessage(message);
        throw e;
      }
    }
  }
}
